package org.example.algorithm.binarytree;

import org.example.algorithm.TreeNode;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

public class VerticalOrderTraversalSolution3 {
    private int minColumn = Integer.MAX_VALUE;

    public List<List<Integer>> verticalTraversal(TreeNode root) {
        List<List<Integer>> res = new ArrayList<>();
        if (root == null) {
            return res;
        }
        Map<Integer, List<int[]>> map = new HashMap<>();
        dfs(root, 0, 0, map);
        for (int i=minColumn;i<minColumn+map.size();i++) {
            List<int[]> list = map.get(i);
            list.sort((a, b) -> a[0] != b[0] ? a[0] - b[0] : a[1] - b[1]);
            List<Integer> newList = new ArrayList<>();
            for (int[] node: list) {
                newList.add(node[1]);
            }
            res.add(newList);
        }
        return res;
    }

    private void dfs(TreeNode root, int row, int column, Map<Integer,List<int[]>> map) {
        minColumn = Math.min(minColumn, column);
        map.computeIfAbsent(column, k->new ArrayList<>()).add(new int[]{row, root.val});
        if (root.left != null) {
            dfs(root.left, row+1, column-1, map);
        }
        if (root.right != null) {
            dfs(root.right, row+1, column+1, map);
        }
    }
}
